/*
 * Copyright (c) 2025 [ByteDance Ltd. and/or its affiliates.] and/or its affiliates
 * SPDX-License-Identifier: LGPL v2.1+
*/
#include "FFmpegBaseXORIOCryptor.h"
#include "FFErrorCodeDef.h"
#include <cassert>
extern "C" {
    #include "libavutil/sha.h"
}

int FFmpegBaseXORIOCryptor::parseParam(const std::map<std::string, FFAny> *params) {
    auto it = params->find(kFFmpegIOCryptorParam_Step);
    if (it == params->end()) {
        LOG_PARAM_NOT_FOUND(kFFmpegIOCryptorParam_Step)
        return TE_ERR_INVALID_PARAM;
    }

    if (!it->second.get(m_step) || 0 == m_step) {
        LOG_PARAM_INVALID(kFFmpegIOCryptorParam_Step)
        return TE_ERR_INVALID_PARAM;
    }

    it = params->find(kFFmpegIOCryptorParam_Length);
    if (it == params->end()) {
        LOG_PARAM_NOT_FOUND(kFFmpegIOCryptorParam_Length)
        return TE_ERR_INVALID_PARAM;
    }

    if (!it->second.get(m_length) || 0 == m_length) {
        LOG_PARAM_INVALID(kFFmpegIOCryptorParam_Length)
        return TE_ERR_INVALID_PARAM;
    }

    it = params->find(kFFmpegIOCryptorParam_Key);
    if (it == params->end()) {
        LOG_PARAM_NOT_FOUND(kFFmpegIOCryptorParam_Key)
        return TE_ERR_INVALID_PARAM;
    }

    if (!it->second.get(m_key)) {
        LOG_PARAM_INVALID(kFFmpegIOCryptorParam_Key)
        return TE_ERR_INVALID_PARAM;
    }

    return 0;
}

int FFmpegBaseXORIOCryptor::readPacket(uint8_t *buf, int bufSize) {
    const auto srcLen = m_pOriginalIOCtx->read_packet(m_pOriginalIOCtx->opaque, buf, bufSize);
    if (srcLen <= 0 || !m_shouldDecrypt)
        return srcLen;

    const int64_t bufPos = m_pCryptIOCtx->pos;
    auto ptr = buf;
    const auto end = buf + srcLen;
    while (ptr < end) {
        const int64_t ptrPos = ptr - buf + bufPos;
        const int64_t multiplesOfStep = ptrPos - ptrPos % m_step;
        const auto nextPtr = ptr + m_step - ptrPos % m_step;   // nextPtr point at position which is multiplies of m_step
        if (ptrPos < multiplesOfStep + m_length) {             // multiplesOfStep <= ptrPos < multiplesOfStep + m_length
            const auto leftBytes = end - ptr;
            int64_t decryptLen = multiplesOfStep + m_length - ptrPos;
            decryptLen = decryptLen < leftBytes ? decryptLen : leftBytes;

            const auto decryptEnd = ptr + decryptLen;
            while (ptr < decryptEnd) {
                *ptr ^= m_key;
                ++ptr;
            }
        }
        ptr = nextPtr;
    }

    return srcLen;
}

int FFmpegBaseXORIOCryptor::writePacket(const uint8_t *buf, int bufSize) {
    const int64_t bufPos = m_pCryptIOCtx->pos;
    auto ptr = const_cast<uint8_t*>(buf);
    const auto end = buf + bufSize;
    while (ptr < end) {
        const int64_t ptrPos = ptr - buf + bufPos;
        const int64_t multiplesOfStep = ptrPos - ptrPos % m_step;
        const auto nextPtr = ptr + m_step - ptrPos % m_step;   // nextPtr point at position which is multiplies of m_step
        if (ptrPos < multiplesOfStep + m_length) {             // multiplesOfStep <= ptrPos < multiplesOfStep + m_length
            const auto leftBytes = end - ptr;
            int64_t decryptLen = multiplesOfStep + m_length - ptrPos;
            decryptLen = decryptLen < leftBytes ? decryptLen : leftBytes;

            const auto decryptEnd = ptr + decryptLen;
            while (ptr < decryptEnd) {
                *ptr ^= m_key;
                ++ptr;
            }
        }
        ptr = nextPtr;
    }

    return m_pOriginalIOCtx->write_packet(m_pOriginalIOCtx->opaque, buf, bufSize);
}

static void readInteger(uint32_t value, uint8_t *dst, uint8_t size) {
    assert(size >= 4);
    for (int i = 0; i < 4; i++) {
        dst[i] = (value >> ((3 - i) * 8)) & 0x000000FF;
    }
}

void FFmpegBaseXORIOCryptor::encryptParamSha256(uint8_t dst[32]) {
    struct AVSHA *sha = av_sha_alloc();
    av_sha_init(sha, 256);

    uint8_t src[9];
    readInteger(m_step, src, 9);
    readInteger(m_length, src + 4, 5);
    src[8] = m_key;

    av_sha_update(sha, src, 9);
    av_sha_final(sha, dst);
    av_free(sha);
}
